import os
import glob
import mne
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# Set your data path
data_path = "/Users/jananjahed/Desktop/Bachelor's project/ds005873"

# -------------------------------
# 1. Load All TSV Annotations
# -------------------------------
def load_all_annotations(tsv_filepaths):
    annotations_list = []
    for fp in tsv_filepaths:
        ann = pd.read_csv(fp, sep='\t')
        # Add the source filename for reference
        ann['source_file'] = os.path.basename(fp)
        annotations_list.append(ann)
    if annotations_list:
        all_annotations = pd.concat(annotations_list, ignore_index=True)
        return all_annotations
    else:
        return pd.DataFrame()  # Return an empty DataFrame if none found

all_tsv_files = glob.glob(os.path.join(data_path, "**", "*_events.tsv"), recursive=True)
print(f"Found {len(all_tsv_files)} TSV files total.")
all_annotations = load_all_annotations(all_tsv_files)

# Count seizure events from annotations (assuming each row is a seizure event)
seizure_count = len(all_annotations)
print(f"Total seizure events in annotations: {seizure_count}")

# ----------------------------------
# 2. Define Epoch Creation Function
# ----------------------------------
def create_epochs(raw, annotations, tmin=-0.2, tmax=0.8):
    events = []
    event_id = {'seizure': 1, 'non-seizure': 0}
    
    # Add seizure events from annotations if available
    if not annotations.empty:
        for _, row in annotations.iterrows():
            onset = row['onset']  # Ensure your TSV has an 'onset' column
            events.append([int(onset * raw.info['sfreq']), 0, event_id['seizure']])
    
    # Add fixed-length non-seizure events (e.g., every 10 seconds)
    non_seizure_events = mne.make_fixed_length_events(raw, id=event_id['non-seizure'], duration=10)
    events.extend(non_seizure_events)
    
    events = np.array(events)
    # Set event_repeated='drop' to handle duplicate event times
    epochs = mne.Epochs(raw, events, event_id, tmin, tmax, baseline=(None, 0), preload=True,
                        event_repeated='drop')
    return epochs


# ----------------------------------
# 3. Process All EDF Files in the Dataset
# ----------------------------------
all_edf_files = glob.glob(os.path.join(data_path, "**", "*_eeg.edf"), recursive=True)
print(f"Found {len(all_edf_files)} EDF files total.")

all_epochs_list = []
for edf_file in all_edf_files:
    print(f"Processing EDF file: {edf_file}")
    
    # Load raw EEG data
    raw = mne.io.read_raw_edf(edf_file, preload=True)
    
    # Plot raw data with all channels (should be 2 channels)
    n_channels_total = len(raw.ch_names)
    raw.plot(n_channels=n_channels_total, duration=10, title=f'Raw EEG Data: {os.path.basename(edf_file)}')
    
    # Preprocess: Filter (1-40 Hz) and Notch Filter (50 Hz)
    raw_filtered = raw.copy().filter(l_freq=1, h_freq=40, method='fir', fir_window='hamming')
    raw_filtered.notch_filter(freqs=[50])
    raw_filtered.plot(n_channels=n_channels_total, duration=10, title=f'Filtered EEG Data: {os.path.basename(edf_file)}')
    
    # Find matching TSV annotation file based on EDF file base name
    base_name = os.path.basename(edf_file).replace('_eeg.edf', '')
    matching_tsv = [fp for fp in all_tsv_files if base_name in fp]
    
    if matching_tsv:
        annotations = pd.read_csv(matching_tsv[0], sep='\t')
        print(f"Using annotations from: {matching_tsv[0]}")
    else:
        print(f"No matching TSV found for {edf_file}. No seizure events will be used for this file.")
        annotations = pd.DataFrame()  # No annotations available
    
    # Create epochs using the filtered data and annotations
    epochs = create_epochs(raw_filtered, annotations)
    print(f"Created {len(epochs.events)} epochs for {os.path.basename(edf_file)}")
    all_epochs_list.append(epochs)

# ----------------------------------
# 4. Combine All Epochs
# ----------------------------------
if all_epochs_list:
    all_epochs = mne.concatenate_epochs(all_epochs_list)
    print("Total epochs created from all files:", len(all_epochs.events))
else:
    print("No epochs were created from the dataset.")

# ----------------------------------
# 5. (Optional) Plot Evoked Response from Combined Epochs
# ----------------------------------
if all_epochs_list:
    evoked = all_epochs.average()
    evoked.plot(spatial_colors=True, time_unit='s', title='Evoked Potentials from Combined Epochs')

Found 2853 TSV files total.
Total seizure events in annotations: 3664
Found 2853 EDF files total.
Processing EDF file: /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-09_eeg.edf
Extracting EDF parameters from /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-09_eeg.edf...
EDF file detected
Setting channel info structure...
Creating raw.info structure...
Reading 0 ... 19222015  =      0.000 ... 75085.996 secs...

Filtering raw data in 1 contiguous segment
Setting up band-pass filter from 1 - 40 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandpass filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 1.00
- Lower transition bandwidth: 1.00 Hz (-6 dB cutoff frequency: 0.50 Hz)
- Upper passband edge: 40.00 Hz
- Upper transition bandwidth: 10.00 Hz (-6 dB cutoff frequency: 45.00 Hz)
- Filter length: 845 samples (3.301 s)

Filtering raw data in 1 contiguous segment
Setting up band-stop filter from 49 - 51 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandstop filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 49.38
- Lower transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 49.12 Hz)
- Upper passband edge: 50.62 Hz
- Upper transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 50.88 Hz)
- Filter length: 1691 samples (6.605 s)

Using annotations from: /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-09_events.tsv
Multiple event values for single event times found. Keeping the first occurrence and dropping all others.
Not setting metadata
7510 matching events found
Setting baseline interval to [-0.19921875, 0.0] s
Applying baseline correction (mode: mean)
0 projection items activated
Using data from preloaded Raw for 7510 events and 257 original time points ...
1 bad epochs dropped
Created 7509 epochs for sub-021_ses-01_task-szMonitoring_run-09_eeg.edf
Processing EDF file: /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-08_eeg.edf
Extracting EDF parameters from /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-08_eeg.edf...
EDF file detected
Setting channel info structure...
Creating raw.info structure...
Reading 0 ... 19253247  =      0.000 ... 75207.996 secs...

Filtering raw data in 1 contiguous segment
Setting up band-pass filter from 1 - 40 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandpass filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 1.00
- Lower transition bandwidth: 1.00 Hz (-6 dB cutoff frequency: 0.50 Hz)
- Upper passband edge: 40.00 Hz
- Upper transition bandwidth: 10.00 Hz (-6 dB cutoff frequency: 45.00 Hz)
- Filter length: 845 samples (3.301 s)

Filtering raw data in 1 contiguous segment
Setting up band-stop filter from 49 - 51 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandstop filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 49.38
- Lower transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 49.12 Hz)
- Upper passband edge: 50.62 Hz
- Upper transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 50.88 Hz)
- Filter length: 1691 samples (6.605 s)

Using annotations from: /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-08_events.tsv
Not setting metadata
7522 matching events found
Setting baseline interval to [-0.19921875, 0.0] s
Applying baseline correction (mode: mean)
0 projection items activated
Using data from preloaded Raw for 7522 events and 257 original time points ...

/var/folders/s6/5tcrcv3n2mb5m18jfn2czp_r0000gn/T/ipykernel_5102/1141197103.py:54: RuntimeWarning: The events passed to the Epochs constructor are not chronologically ordered.
  epochs = mne.Epochs(raw, events, event_id, tmin, tmax, baseline=(None, 0), preload=True,

1 bad epochs dropped
Created 7521 epochs for sub-021_ses-01_task-szMonitoring_run-08_eeg.edf
Processing EDF file: /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-01_eeg.edf
Extracting EDF parameters from /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-01_eeg.edf...
EDF file detected
Setting channel info structure...
Creating raw.info structure...
Reading 0 ... 3504383  =      0.000 ... 13688.996 secs...

Filtering raw data in 1 contiguous segment
Setting up band-pass filter from 1 - 40 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandpass filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 1.00
- Lower transition bandwidth: 1.00 Hz (-6 dB cutoff frequency: 0.50 Hz)
- Upper passband edge: 40.00 Hz
- Upper transition bandwidth: 10.00 Hz (-6 dB cutoff frequency: 45.00 Hz)
- Filter length: 845 samples (3.301 s)

Filtering raw data in 1 contiguous segment
Setting up band-stop filter from 49 - 51 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandstop filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 49.38
- Lower transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 49.12 Hz)
- Upper passband edge: 50.62 Hz
- Upper transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 50.88 Hz)
- Filter length: 1691 samples (6.605 s)

Using annotations from: /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-01_events.tsv
Not setting metadata
1369 matching events found
Setting baseline interval to [-0.19921875, 0.0] s
Applying baseline correction (mode: mean)
0 projection items activated
Using data from preloaded Raw for 1369 events and 257 original time points ...
1 bad epochs dropped
Created 1368 epochs for sub-021_ses-01_task-szMonitoring_run-01_eeg.edf
Processing EDF file: /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-03_eeg.edf
Extracting EDF parameters from /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-03_eeg.edf...
EDF file detected
Setting channel info structure...
Creating raw.info structure...
Reading 0 ... 2582527  =      0.000 ... 10087.996 secs...

/var/folders/s6/5tcrcv3n2mb5m18jfn2czp_r0000gn/T/ipykernel_5102/1141197103.py:54: RuntimeWarning: The events passed to the Epochs constructor are not chronologically ordered.
  epochs = mne.Epochs(raw, events, event_id, tmin, tmax, baseline=(None, 0), preload=True,

raw_filtered = raw.copy()
raw_filtered.filter(l_freq=1, h_freq=40, method='fir', fir_window='hamming')

Filtering raw data in 1 contiguous segment
Setting up band-pass filter from 1 - 40 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandpass filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 1.00
- Lower transition bandwidth: 1.00 Hz (-6 dB cutoff frequency: 0.50 Hz)
- Upper passband edge: 40.00 Hz
- Upper transition bandwidth: 10.00 Hz (-6 dB cutoff frequency: 45.00 Hz)
- Filter length: 845 samples (3.301 s)

raw_filtered.notch_filter(freqs=[50])
raw_filtered.plot(n_channels=2, duration=10, title='Filtered EEG Data (1-40 Hz, Notch 50 Hz)')

Filtering raw data in 1 contiguous segment
Setting up band-stop filter from 49 - 51 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandstop filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 49.38
- Lower transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 49.12 Hz)
- Upper passband edge: 50.62 Hz
- Upper transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 50.88 Hz)
- Filter length: 1691 samples (6.605 s)

def load_all_annotations(filepaths):
    annotations_list = []
    for fp in filepaths:
        ann = pd.read_csv(fp, sep='\t')
        ann['source_file'] = os.path.basename(fp)
        annotations_list.append(ann)
    if annotations_list:
        return pd.concat(annotations_list, ignore_index=True)
    else:
        return pd.DataFrame()  # empty if no files found

all_tsv_files = glob.glob(os.path.join(data_path, "**", "*_events.tsv"), recursive=True)
print(f"Found {len(all_tsv_files)} TSV files total.")
annotations_df = load_all_annotations(all_tsv_files)

# --- Count Seizure and Non-Seizure Events ---
# Here we assume that every row in your TSV is a seizure event.
# If your file includes a column (e.g., "event" or "label") to distinguish, adjust accordingly.
seizure_count = len(annotations_df)
print(f"Seizure events (from annotations): {seizure_count}")

# If your TSV has a column indicating event type, for example:
# seizure_count = (annotations_df['event'] == 'seizure').sum()
# non_seizure_count = (annotations_df['event'] == 'non-seizure').sum()
# print(f"Seizure events: {seizure_count}, Non-seizure events: {non_seizure_count}")

# --- Create Epochs for a Given Raw Data and Its Annotations ---
def create_epochs(raw, annotations, tmin=-0.2, tmax=0.8):
    events = []
    event_id = {'seizure': 1, 'non-seizure': 0}
    
    # Add seizure events from the annotations
    for _, row in annotations.iterrows():
        onset = row['onset']
        events.append([int(onset * raw.info['sfreq']), 0, event_id['seizure']])
    
    # Add non-seizure events with fixed-length epochs (e.g., every 10 seconds)
    non_seizure_events = mne.make_fixed_length_events(raw, id=event_id['non-seizure'], duration=10)
    events.extend(non_seizure_events)
    
    events = np.array(events)
    epochs = mne.Epochs(raw, events, event_id, tmin, tmax, baseline=(None, 0), preload=True)
    return epochs

# --- Process All EDF Files in the Dataset ---
all_edf_files = glob.glob(os.path.join(data_path, "**", "*_eeg.edf"), recursive=True)
print(f"Found {len(all_edf_files)} EDF files total.")

epochs_list = []
for edf_file in all_edf_files:
    # Load raw EEG data
    raw = mne.io.read_raw_edf(edf_file, preload=True)
    
    # Optionally, preprocess your raw data here (e.g., filtering, notch filtering, etc.)
    raw.filter(l_freq=1, h_freq=40, method='fir', fir_window='hamming')
    raw.notch_filter(freqs=[50])
    
    # Find the corresponding TSV file using part of the file name
    base_name = os.path.basename(edf_file).replace('_eeg.edf', '')
    matching_tsv = [fp for fp in all_tsv_files if base_name in fp]
    
    if matching_tsv:
        annotations = pd.read_csv(matching_tsv[0], sep='\t')
    else:
        print(f"No matching TSV found for {edf_file}. Skipping seizure events for this file.")
        annotations = pd.DataFrame()  # No seizure events available
    
    epochs = create_epochs(raw, annotations)
    epochs_list.append(epochs)

# Optionally, combine epochs from all files into one object (if they have the same channel layout)
if epochs_list:
    all_epochs = mne.concatenate_epochs(epochs_list)
    print("Total epochs created:", len(all_epochs.events))
else:
    print("No epochs were created.")

Found 2853 TSV files total.
Seizure events (from annotations): 3664
Found 2853 EDF files total.
Extracting EDF parameters from /Users/jananjahed/Desktop/Bachelor's project/ds005873/sub-021/ses-01/eeg/sub-021_ses-01_task-szMonitoring_run-09_eeg.edf...
EDF file detected
Setting channel info structure...
Creating raw.info structure...
Reading 0 ... 19222015  =      0.000 ... 75085.996 secs...
Filtering raw data in 1 contiguous segment
Setting up band-pass filter from 1 - 40 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandpass filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 1.00
- Lower transition bandwidth: 1.00 Hz (-6 dB cutoff frequency: 0.50 Hz)
- Upper passband edge: 40.00 Hz
- Upper transition bandwidth: 10.00 Hz (-6 dB cutoff frequency: 45.00 Hz)
- Filter length: 845 samples (3.301 s)

Filtering raw data in 1 contiguous segment
Setting up band-stop filter from 49 - 51 Hz

FIR filter parameters
---------------------
Designing a one-pass, zero-phase, non-causal bandstop filter:
- Windowed time-domain design (firwin) method
- Hamming window with 0.0194 passband ripple and 53 dB stopband attenuation
- Lower passband edge: 49.38
- Lower transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 49.12 Hz)
- Upper passband edge: 50.62 Hz
- Upper transition bandwidth: 0.50 Hz (-6 dB cutoff frequency: 50.88 Hz)
- Filter length: 1691 samples (6.605 s)

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[19], line 68
     65         print(f"No matching TSV found for {edf_file}. Skipping seizure events for this file.")
     66         annotations = pd.DataFrame()  # No seizure events available
---> 68     epochs = create_epochs(raw, annotations)
     69     epochs_list.append(epochs)
     71 # Optionally, combine epochs from all files into one object (if they have the same channel layout)

Cell In[19], line 42, in create_epochs(raw, annotations, tmin, tmax)
     39 events.extend(non_seizure_events)
     41 events = np.array(events)
---> 42 epochs = mne.Epochs(raw, events, event_id, tmin, tmax, baseline=(None, 0), preload=True)
     43 return epochs

File <decorator-gen-228>:12, in __init__(self, raw, events, event_id, tmin, tmax, baseline, picks, preload, reject, flat, proj, decim, reject_tmin, reject_tmax, detrend, on_missing, reject_by_annotation, metadata, event_repeated, verbose)

File /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/mne/epochs.py:3580, in Epochs.__init__(self, raw, events, event_id, tmin, tmax, baseline, picks, preload, reject, flat, proj, decim, reject_tmin, reject_tmax, detrend, on_missing, reject_by_annotation, metadata, event_repeated, verbose)
   3575     events, event_id, annotations = _events_from_annotations(
   3576         raw, events, event_id, annotations, on_missing
   3577     )
   3579 # call BaseEpochs constructor
-> 3580 super().__init__(
   3581     info,
   3582     None,
   3583     events,
   3584     event_id,
   3585     tmin,
   3586     tmax,
   3587     metadata=metadata,
   3588     baseline=baseline,
   3589     raw=raw,
   3590     picks=picks,
   3591     reject=reject,
   3592     flat=flat,
   3593     decim=decim,
   3594     reject_tmin=reject_tmin,
   3595     reject_tmax=reject_tmax,
   3596     detrend=detrend,
   3597     proj=proj,
   3598     on_missing=on_missing,
   3599     preload_at_end=preload,
   3600     event_repeated=event_repeated,
   3601     verbose=verbose,
   3602     raw_sfreq=raw_sfreq,
   3603     annotations=annotations,
   3604 )

File <decorator-gen-212>:12, in __init__(self, info, data, events, event_id, tmin, tmax, baseline, raw, picks, reject, flat, decim, reject_tmin, reject_tmax, detrend, proj, on_missing, preload_at_end, selection, drop_log, filename, metadata, event_repeated, raw_sfreq, annotations, verbose)

File /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/mne/epochs.py:529, in BaseEpochs.__init__(***failed resolving arguments***)
    520     self.drop_log = drop_log
    522 self.events = self.events[selected]
    524 (
    525     self.events,
    526     self.event_id,
    527     self.selection,
    528     self.drop_log,
--> 529 ) = _handle_event_repeated(
    530     self.events,
    531     self.event_id,
    532     event_repeated,
    533     self.selection,
    534     self.drop_log,
    535 )
    537 # then subselect
    538 sub = np.where(np.isin(selection, self.selection))[0]

File /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/mne/epochs.py:330, in _handle_event_repeated(events, event_id, event_repeated, selection, drop_log)
    328 drop_log = list(drop_log)
    329 if event_repeated == "error":
--> 330     raise RuntimeError(
    331         "Event time samples were not unique. Consider "
    332         'setting the `event_repeated` parameter."'
    333     )
    335 elif event_repeated == "drop":
    336     logger.info(
    337         "Multiple event values for single event times found. "
    338         "Keeping the first occurrence and dropping all others."
    339     )

RuntimeError: Event time samples were not unique. Consider setting the `event_repeated` parameter."

	General
	Filename(s)	sub-021_ses-01_task-szMonitoring_run-09_eeg.edf
	MNE object type	RawEDF
	Measurement date	2000-01-01 at 00:00:00 UTC
	Participant	x
	Experimenter	Unknown
	Acquisition
	Duration	20:51:26 (HH:MM:SS)
	Sampling frequency	256.00 Hz
	Time points	19,222,016
	Channels
	EEG
	Head & sensor digitization	Not available
	Filters
	Highpass	1.00 Hz
	Lowpass	40.00 Hz